package edu.kit.csl.cratyli.control;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.DataInputStream;
import java.io.FileInputStream;
import java.io.FileWriter;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;

import edu.kit.csl.cratyli.datatypes.BadParameterValueException;
import edu.kit.csl.cratyli.datatypes.Dictionary;
import edu.kit.csl.cratyli.datatypes.TextCorpus;
import edu.kit.csl.cratyli.datatypes.Word;
import edu.kit.csl.cratyli.ui.Configuration;
import edu.kit.csl.cratyli.ui.Logger;

/**
 * This combination strategy runs the m2m-aligner (available at
 * {@linkplain https://code.google.com/p/m2m-aligner/}) to align the primary
 * and secondary text corpus. The combination is done by merging the 1:1
 * relations, inserting 1:0 relations, taking the most granular option in 
 * 1:n relations, and the primary path in n:m relations. The approach is
 * described in (Stahlberg et. al., 2013). 
 */
public class M2MCombinationStrategy extends CombinationStrategy {
	
	private final String TMP_FILE = "/tmp/cratyli.tmp.m2m."
			+ System.currentTimeMillis() + ".";
	private final String m2mAlignerCmd;
	
	private Set<String> toMerge;
	
	private final String[] envStrings;
	private final double cutOff, secWeightLimit;
	
	private SentenceTargetFunction senTrgtFunction;

	/**
	 * Sole constructor. Fetches the m2mAlignerCommand parameter
	 */
	public M2MCombinationStrategy() throws BadParameterValueException {
		super();
		Configuration conf = Configuration.getSingleton();
		m2mAlignerCmd = conf.getString("m2mAlignerCommand") 
			+ " -i " + TMP_FILE + "toAlign"
			+ " -o " + TMP_FILE + "out --alignerOut " + TMP_FILE
			+ "model --errorInFile";
		Logger.getSingleton().debug("M2M Aligner command: " + m2mAlignerCmd);
		toMerge = new HashSet<String>();
		final Map<String, String> env = 
				new HashMap<String, String>(System.getenv());
        env.put("LD_LIBRARY_PATH", conf.getString("m2mLibraryPath"));
        envStrings = mapToStringArray(env);
        cutOff = conf.getFloat("m2mCutOffProbability");
        secWeightLimit = conf.getFloat("secondaryWeightLimit");
        String f = conf.getString("sentenceTargetFunction");
		if (f.equals("mostGranular")) {
			senTrgtFunction = new LengthSentenceTargetFunction();
		} else if (f.equals("weightSum")) {
			senTrgtFunction = new WeightSentenceTargetFunction();
		}  else if (f.equals("weightMin")) {
			senTrgtFunction = new MinWeightSentenceTargetFunction();
		} else {
			throw new BadParameterValueException(
					"Sentence target function not known");
		}
	}
	
	/* (non-Javadoc)
	 * @see CombinationStrategy#combine(Dictionary, TextCorpus, TextCorpus)
	 */
	@Override
	public TextCorpus combine(Dictionary dict, TextCorpus primary,
			TextCorpus secondary) {
		toMerge.clear();
		merger.next();
		try {
			FileWriter fs;
			fs = new FileWriter(TMP_FILE + "toAlign");
			BufferedWriter out = new BufferedWriter(fs);
			ArrayList<String> ids = new ArrayList<String>();
			Map<String, String> additional = new HashMap<String, String>();
			for (Entry<String, String> entry : primary) {
				String id = entry.getKey();
				String primSen = entry.getValue();
				String secSen = secondary.get(id);
				if (secSen == null || secSen.isEmpty()) {
					additional.put(id, primSen);
				} else if (primSen.isEmpty()) {
					additional.put(id, secSen);
				} else {
					ids.add(id);
					out.write(entry.getValue() + "\t" + secSen + "\n");
				}
			}
			out.close();
			runM2MAligner();
			TextCorpus corpus = createCorpusFromM2MOutput(dict, primary, ids);
			corpus.putAll(additional);
			mergePronunciations(dict);
			return corpus;
		} catch (IOException e) {
			Logger.getSingleton().fatalError("Could not write to m2m alignment"
				+ " file " + TMP_FILE + "toAlign: " + e.getMessage());
		} catch (InterruptedException e) {
			Logger.getSingleton().error("Program execution of m2m aligner was "
					+ "interrupted: " + e.getMessage());
		} catch (NumberFormatException e) {
			Logger.getSingleton().error("M2M alignment model file contains a "
				+ "malformed probability: " + e.getMessage());
		}
		return new TextCorpus();
	}
	
	/*
	 * Create the merged output from the output of the M2M Aligner output. The
	 * inner for loop implements the approach presented in 
	 * (Stahlberg et. al., 2013) to deal with x:y alignments.
	 */
	private TextCorpus createCorpusFromM2MOutput(Dictionary dict,
			TextCorpus primCorpus, ArrayList<String> ids) throws IOException {
		FileInputStream fstream = new FileInputStream(TMP_FILE + "out");
		DataInputStream in = new DataInputStream(fstream);
		BufferedReader br = new BufferedReader(new InputStreamReader(in));
		String line;
		TextCorpus corpus = new TextCorpus();
		for (int lineNum = 0; (line = br.readLine()) != null; lineNum++) {
			line = line.trim();
			String id = ids.get(lineNum);
			if (line.equals("NO ALIGNMENT")) {
				corpus.put(id, primCorpus.get(id));
				continue;
			}
			try { // TODO: Bad practice: using try for control flow..
				String[] splitLine = line.split("\t", 2);
				String[] pSplit = splitLine[0].split("\\|");
				String[] sSplit = splitLine[1].split("\\|");
				StringBuilder newSen = new StringBuilder();
				if (pSplit.length != sSplit.length) {
					Logger.getSingleton().error("Number of 1st level "
						+ "separators does not match in " + line);
					continue;
				}
				for (int i = 0; i < pSplit.length; i++) {
					handleAlignment(dict, newSen, pSplit[i], sSplit[i]);
				}
				corpus.put(id, newSen.toString().trim());
			} catch (ArrayIndexOutOfBoundsException e) {
				Logger.getSingleton().error("Malformed line in m2m aligner "
					+ "output: " + line);
			}
		}
		in.close();
		return corpus;
	}
	
	/*
	 * Handle x:y alignments (x,y>0) and update the new sentence. Take the most
	 * granular path in n:m alignments, register 1:1 for merging the 
	 * pronunciations and take the primary path.
	 */
	private void handleAlignment(Dictionary dict, StringBuilder newSen,
			String pAlign, String sAlign) {
		if (pAlign.equals("_")) { // 0:1
			newSen.append(" " + sAlign.replaceAll(":", " "));
		} else if (sAlign.equals("_")) { // 1:0
			newSen.append(" " + pAlign.replaceAll(":", " "));
		} else { // 1:1, 1:n, n:1, n:m
			String[] pWords = pAlign.split(":");
			String[] sWords = sAlign.split(":");
			if (pWords.length == 1 && sWords.length == 1) { // 1:1
				newSen.append(" " + pWords[0]);
				toMerge.add(pWords[0]);
			} else {
				newSen.append(" " 
						+ senTrgtFunction.merge(dict, pWords, sWords));
			}
		}
	}
	
	/*
	 * Merges the pronunciations according the m2m alignment model.
	 */
	private void mergePronunciations(Dictionary dict) throws IOException,
			NumberFormatException {
		Logger log = Logger.getSingleton();
		FileInputStream fstream = new FileInputStream(TMP_FILE + "model");
		DataInputStream in = new DataInputStream(fstream);
		BufferedReader br = new BufferedReader(new InputStreamReader(in));
		String lastWord = null, line;
		Word pWord = null, sWord = null;
		Map<Word, Double> sWords = new HashMap<Word, Double>();
		double score;
		while ((line = br.readLine()) != null) {
			String[] entry = line.trim().split("\\s+", 3);
			if (entry.length != 3) {
				log.warn("Invalid line '" + line + "' in m2m model file");
				continue;
			}
			String thisWord = entry[0].trim();
			if (!thisWord.equals(lastWord)) {
				if (pWord != null && !sWords.isEmpty()) {
					merger.merge(pWord, sWords);
					sWords.clear();
				}
				pWord = toMerge.contains(thisWord) ? dict.get(thisWord) : null;
				lastWord = thisWord;
			}
			if (pWord != null && (sWord = dict.get(entry[1].trim())) != null && 
					sWord.getWeight() <= pWord.getWeight() * secWeightLimit &&
					(score = Double.parseDouble(entry[2].trim())) >= cutOff) {
				sWords.put(sWord, score * sWord.getWeight());
			}
		}
		if (pWord != null && !sWords.isEmpty()) {
			merger.merge(pWord, sWords);
		}
		in.close();
	}
	
	
	/*
	 * Run the M2M Aligner.
	 */
	private void runM2MAligner() throws IOException, InterruptedException  {
		Logger log = Logger.getSingleton();
		log.notice("Run m2m Aligner");
		Process p = Runtime.getRuntime().exec(m2mAlignerCmd, envStrings);
		BufferedReader bri = new BufferedReader
				(new InputStreamReader(p.getInputStream()));
		BufferedReader bre = new BufferedReader
				(new InputStreamReader(p.getErrorStream()));
		String line;
		while ((line = bri.readLine()) != null) {
			log.notice("m2m-aligner stdout: " + line);
		}
		while ((line = bre.readLine()) != null) {
			log.error("m2m-aligner stderr: " + line);
		}
		bri.close();
		bre.close();
		p.waitFor();
	}
	
	/*
	 * @see http://stackoverflow.com/questions/2741175/set-windows-path-environment-variable-at-runtime-in-java
	 */
    private static String[] mapToStringArray(Map<String, String> map) {
        final String[] strings = new String[map.size()];
        int i = 0;
        for (Map.Entry<String, String> e : map.entrySet()) {
            strings[i] = e.getKey() + '=' + e.getValue();
            i++;
        }
        return strings;
    }
}
